package com.csust.automaticteach.app.advisor;

import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Objects;

/**
 * @author yuluo
 * @author <a href="mailto:yuluo08290126@gmail.com">yuluo</a>
 * Incorporate DeepSeek-R1's reasoning content into the output
 */

public class ReasoningContentAdvisor implements BaseAdvisor {


    private final int order;

    public ReasoningContentAdvisor(Integer order) {
        this.order = order != null ? order : 0;
    }

    @Override
    public int getOrder() {

        return this.order;
    }

    @Override
    public AdvisedRequest before(AdvisedRequest request) {
        return request;
    }

    @Override
    public AdvisedResponse after(AdvisedResponse advisedResponse) {
        //根据 response 得到 ChatResponse
        ChatResponse resp = advisedResponse.response();
        if (Objects.isNull(resp)) {
            return advisedResponse;
        }
        String reasoningContent = String.valueOf(resp.getResults().get(0).getOutput().getMetadata().get("reasoningContent"));
        if (StringUtils.hasText(reasoningContent)) {
            List<Generation> thinkGenerations = resp.getResults().stream()
                    .map(generation -> {
                        AssistantMessage output = generation.getOutput();
                        AssistantMessage thinkAssistantMessage = new AssistantMessage(
                                String.format("<think>%s</think>", reasoningContent) + output.getText(),
                                output.getMetadata(),
                                output.getToolCalls(),
                                output.getMedia()
                        );
                        return new Generation(thinkAssistantMessage, generation.getMetadata());
                    }).toList();

            ChatResponse thinkChatResp = ChatResponse.builder().from(resp).generations(thinkGenerations).build();
            return AdvisedResponse.builder().adviseContext(advisedResponse.adviseContext()).response(thinkChatResp).build();

        }
        return advisedResponse;
    }

    @Override
    public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
        return BaseAdvisor.super.aroundCall(advisedRequest, chain);
    }

    @Override
    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
        return BaseAdvisor.super.aroundStream(advisedRequest, chain);
    }
}